namespace SimpleTPL;

public class SimpleTask<TResult> : SimpleTask
{
    private Func<TResult> _function;
    private TResult _result;

    public TResult Result
    {
        get
        {
            if (!IsCompleted) throw new InvalidOperationException("Task not completed");
            if (_exception != null) throw _exception;
            return _result;
        }
    }
    
    public SimpleTaskAwaiter<TResult> GetAwaiter()
    {
        return new SimpleTaskAwaiter<TResult>(this);
    }
    
    public static SimpleAsyncMethodBuilder<TResult> CreateAsyncMethodBuilder() => 
        SimpleAsyncMethodBuilder<TResult>.Create();

    internal new void Execute()
    {
        try
        {
            _result = _function(); // 现在直接使用存储的函数
        }
        catch (Exception ex)
        {
            _exception = ex;
        }
        finally
        {
            lock (_syncLock)
            {
                _isCompleted = true;
                foreach (var continuation in _continuations)
                {
                    try
                    {
                        continuation(this);
                    }
                    catch { /* 忽略延续异常 */ }
                }
                _continuations = null;
            }
        }
    }
    
    internal void SetResult(TResult result)
    {
        lock (_syncLock)
        {
            _result = result;
            _isCompleted = true;
            foreach (var continuation in _continuations)
            {
                try
                {
                    continuation(this);
                }
                catch { /* 忽略延续异常 */ }
            }
            _continuations = null;
        }
    }

    internal void SetException(Exception exception)
    {
        lock (_syncLock)
        {
            _exception = exception;
            _isCompleted = true;
            foreach (var continuation in _continuations)
            {
                try
                {
                    continuation(this);
                }
                catch { /* 忽略延续异常 */ }
            }
            _continuations = null;
        }
    }
    
    public SimpleTask(Func<TResult> function) : base()
    {
        _function = function ?? throw new ArgumentNullException(nameof(function));
        _action = () => _result = _function();
    }

    // 实例方法
    public SimpleTask ContinueWith(Action<SimpleTask<TResult>> continuation)
    {
        return base.ContinueWith(t => continuation((SimpleTask<TResult>)t));
    }

    public SimpleTask<TNewResult> ContinueWith<TNewResult>(Func<SimpleTask<TResult>, TNewResult> continuation)
    {
        var tcs = new SimpleTaskCompletionSource<TNewResult>();

        lock (_syncLock)
        {
            if (IsCompleted)
            {
                SimpleTask.Run(() => tcs.SetResult(continuation(this)));
            }
            else
            {
                base.ContinueWith(t => 
                    SimpleTask.Run(() => tcs.SetResult(continuation((SimpleTask<TResult>)t))));
            }
        }

        return tcs.Task;
    }
}